import os

import cv2
from option import parse_args


def get_max_width(actor_num, frame, video_path):
    w_max = 0
    for i in range(actor_num):
        image_name = '{}_{}.png'.format(frame, i)
        if os.path.exists(os.path.join(video_path, image_name)):
            img = cv2.imread(os.path.join(video_path, image_name), cv2.IMREAD_GRAYSCALE)
            w_max = max(w_max, img.shape[1])
    return w_max


def remove_actors(actor_num, frame, video_path, max_width):
    for i in range(actor_num):
        image_name = '{}_{}.png'.format(frame, i)
        if os.path.exists(os.path.join(video_path, image_name)):
            img = cv2.imread(os.path.join(video_path, image_name), cv2.IMREAD_GRAYSCALE)
            if img.shape[1] < max_width:
                os.remove(os.path.join(video_path, image_name))


def main():
    args = parse_args()
    crops = os.path.join(args.root_dir, 'crops')
    lines = open('multiple_actors_statistics.dat').readlines()
    for line in lines:
        line_split = line.split(',')
        assert len(line_split) == 2
        video_path = os.path.join(crops, line_split[0])
        actor_num = int(line_split[1].strip())

        for crop in os.listdir(video_path):
            crop_split = crop[:-4].split('_')
            if int(crop_split[1]) != 0:
                continue
            max_width = get_max_width(actor_num, crop_split[0], video_path)
            remove_actors(actor_num, crop_split[0], video_path, max_width)


if __name__ == '__main__':
    main()
